import os
import shutil
import re
from pathlib import Path
from tqdm import tqdm
from dicom2nifti import convert_directory

data_root = r"F:\CT ALL"
target_root = r"\\mega\dxli\Data\esophageal_cancer_raw2"
case_list = [os.path.join(data_root,i) for i in os.listdir(data_root)]
founded_data = []
for i in tqdm(case_list):
    case_name =Path(i).name
    # names = str(case_name).split('--')
    target_folder = Path(target_root)/'CT'/case_name
    os.makedirs(target_folder, exist_ok=True)
    if len(os.listdir(target_folder))>0:
        continue
    folders = os.listdir(i)
    for f in folders:
        path = os.path.join(i,f)
        if 3<=len(os.listdir(path))<=5:
            src_data = path
            founded_data.append(src_data)
                
    if len(os.listdir(src_data)) <5:
        for k in os.listdir(src_data):
            data = os.path.join(src_data,k)            
            if os.path.isdir(data) and 'SERIES' not in data:                
                convert_directory(data, target_folder, reorient=False)
                for i in Path(target_folder).iterdir():
                    os.rename(i, str(Path(target_folder)/'image.nii.gz'))
            elif os.path.isfile(data) and 'NII' in k:
                try:
                    shutil.copy(data, target_folder/'label.nii.gz')
                except:
                    print(data,'ROI')

    elif len(os.listdir(src_data)) == 5:
        for k in os.listdir(src_data):
            data = os.path.join(src_data,k)
            if (os.path.isdir(data) and 
                len(str(Path(data).name))>3
                and bool(re.search(r'\d', Path(data).name))):
                os.makedirs(target_folder, exist_ok=True)
                for file in Path(data).iterdir():
                    if '_SRC.NII.GZ' in str(file):
                        try:
                            shutil.copy(file, str(Path(target_folder)/'image.nii.gz'))
                        except:
                            print(src_data,'image')
                    elif '_I.NII.GZ':
                        try:
                            shutil.copy(file, str(Path(target_folder)/'label.nii.gz'))
                        except:
                            print(src_data, 'label')
                    else:
                        print(src_data,'others')
print(len(founded_data))
